package net.hasor.dataql.fx.basic;

import com.alibaba.fastjson.JSONObject;
import net.hasor.dataql.UdfSourceAssembly;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.inject.Singleton;
import javax.sql.DataSource;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.ResultSetMetaData;
import java.util.*;

/**
 * 自定义 Http 工具包
 *
 * @author tdx.lq
 * @since 2020/9/30 7:08
 */

@Singleton
public class SqlUdfSource implements UdfSourceAssembly {
    private static final Logger logger = LoggerFactory.getLogger(SqlUdfSource.class);
    public static Map<String, DataSource> dataSourceMap = new HashMap<>();


    /**
     * 获取连接池
     *
     * @param pool 连接池名
     * @return 数据库连接池
     */
    private Connection getConnectionFromPool(String pool) {
        Connection con;
        try {
            con = dataSourceMap.get(pool).getConnection();
        } catch (Exception e) {
            throw new RuntimeException("获取数据库连接池异常，conn=" + pool, e);
        }
        return con;
    }


    /**
     * 执行sql
     *
     * @param pool   数据库连接池名
     * @param sql    sql语句
     * @param params 参数，对应sql语句里的 “?”
     * @return 受影响条数
     */
    public int update(String pool, String sql, Object... params) {
        Connection conn = null;
        PreparedStatement stmt = null;
        try {
            conn = getConnectionFromPool(pool);
            conn.setAutoCommit(false); // 禁用自动事务提交
            stmt = conn.prepareStatement(sql);
            if (params != null && params.length > 0) {
                for (int i = 0; i < params.length; i++) {
                    stmt.setObject(i + 1, params[i]);
                }
            }
            int r = stmt.executeUpdate();
            conn.commit(); // 手动提交
            return r;
        } catch (Exception e) {
            if (conn != null) {
                try {
                    conn.rollback();
                } catch (Exception e2) {
                    logger.warn("数据库回滚异常，pool={},sql={}", pool, sql, e2);
                }
            }
            if (stmt != null) {
                try {
                    stmt.close();
                } catch (Exception e2) {
                    logger.warn("关闭 stmt 异常，pool={},sql={}", pool, sql, e2);
                }
            }
            if (conn != null) {
                try {
                    conn.close();
                } catch (Exception e2) {
                    logger.warn("关闭连接池异常，pool={},sql={}", pool, sql, e2);
                }
            }
            throw new RuntimeException("数据库执行出错：" + pool + "|" + sql, e);
        }
    }

    /**
     * 获取列表的 JSON 字符串表示
     *
     * @param pool   连接池名
     * @param sql    SQL 语句
     * @param params 参数，对应 SQL 中的 “？”
     * @return 查询结果集的 JSON 字符串表示
     */
    public String strArr(String pool, String sql, Object... params) {
        return JSONObject.toJSONString(query(pool, sql, params));
    }

    /**
     * 获取 MAP 的 JSON 字符串表示
     *
     * @param pool   连接池名
     * @param sql    SQL 语句
     * @param params 参数，对应 SQL 中的 “？”
     * @return 查询结果集的 JSON 字符串表示
     */
    public String strMap(String pool, String sql, Object... params) {
        return JSONObject.toJSONString(Objects.requireNonNull(query(pool, sql, params)).get(0));
    }

    /**
     * 获取结果集的第一行第一列的值
     *
     * @param pool   连接池名
     * @param sql    SQL 语句
     * @param params 参数，对应 SQL 中的 “？”
     * @return 返回第一行第一列的值
     */
    public Object obj(String pool, String sql, Object... params) {
        return Objects.requireNonNull(query(pool, sql, params))
                .get(0)
                .values()
                .iterator()
                .next();
    }

    /**
     * 查询数据库
     *
     * @param pool   连接池名
     * @param sql    查询语句
     * @param params 参数，对应sql语句中的“?”
     * @return 数据集
     */
    public List<LinkedHashMap<String, Object>> query(String pool, String sql, Object... params) {
        try (Connection conn = getConnectionFromPool(pool);
             PreparedStatement stmt = conn.prepareStatement(sql)
        ) {
            if (params != null && params.length > 0)
                for (int i = 0; i < params.length; i++)
                    stmt.setObject(i + 1, params[i]);

            stmt.execute();

            try (ResultSet rs = stmt.getResultSet()) {
                ResultSetMetaData meta = rs.getMetaData();
                int nColumnCount = meta.getColumnCount();
                if (nColumnCount <= 0) {
                    return null;
                }
                int i;
                ArrayList<LinkedHashMap<String, Object>> list = new ArrayList<>();
                while (rs.next()) {
                    LinkedHashMap<String, Object> rows = new LinkedHashMap<>();
                    for (i = 1; i <= nColumnCount; i++) {
                        String colName = meta.getColumnName(i);
                        rows.put(colName, rs.getObject(i));
                    }
                    list.add(rows);
                }
                return list;
            }
        } catch (Exception e) {
            throw new RuntimeException("查询数据库出错：" + pool + "|" + sql, e);
        }
    }
}
